package com.yohann.algorithm;

import java.io.*;
import java.util.*;

/**
 * <p>
 * 哈夫曼编码演示
 * </p>
 *
 * @author Yohann
 * @since 2021/1/6 21:07
 */
public class HuffmanCodeDemo {
    public static void main(String[] args) {
//        String str = "i like like like java do you like a java";
//        byte[] bytes = str.getBytes();
//        byte[] zip = huffmanZip(bytes);
//        byte[] decode = decode(huffmanCodes, zip);
//        System.out.println(new String(decode));

        String srcFile = "D:\\eclipse\\test.bmp";
        String dstFile = "D:\\eclipse\\test.zip";
        String zipFile = "D:\\eclipse\\aaa.bmp";
//        zipFile(srcFile, dstFile);
        decodeFile(dstFile, zipFile);
    }

    public static byte[] huffmanZip(byte[] bytes) {
        List<HuffmanNode> nodes = createNodes(bytes);
        HuffmanNode node = generateHuffmanTree(nodes);
        getHuffmanCodeTable(node, "");
        byte[] zip = zip(bytes);
        return zip;
    }

    public static void zipFile(String srcFile, String dsFile) {
        FileInputStream in = null;
        FileOutputStream out = null;
        ObjectOutputStream oos = null;

        try {
            in = new FileInputStream(srcFile);
            byte[] bytes = new byte[in.available()];
            in.read(bytes);
            byte[] zip = huffmanZip(bytes);

            out = new FileOutputStream(dsFile);

            oos = new ObjectOutputStream(out);
            oos.writeObject(zip);
            oos.writeObject(huffmanCodes);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (in != null) {
                try {
                    in.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (out != null) {
                try {
                    in.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (oos != null) {
                try {
                    in.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    public static void decodeFile(String zipFile, String dstFile) {
        FileInputStream in = null;
        FileOutputStream out = null;
        ObjectInputStream ois = null;

        try {
            in = new FileInputStream(zipFile);
            byte[] bytes = new byte[in.available()];
            ois = new ObjectInputStream(in);

            byte[] huffmanBytes = (byte[]) ois.readObject();
            Map<Byte, String> codeMap = (Map<Byte, String>) ois.readObject();
            byte[] decode = decode(codeMap, huffmanBytes);

            out = new FileOutputStream(dstFile);
            out.write(decode);
        } catch (IOException | ClassNotFoundException e) {
            e.printStackTrace();
        } finally {
            if (out != null) {
                try {
                    in.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (ois != null) {
                try {
                    in.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (in != null) {
                try {
                    in.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    public static List<HuffmanNode> createNodes(byte[] bytes) {
        List<HuffmanNode> huffmanNodes = new ArrayList<>();
        Map<Byte, Integer> countMap = new HashMap<>();

        for (byte b : bytes) {
            Integer count = countMap.get(b);
            if (count == null) {
                countMap.put(b, 1);
            } else {
                countMap.put(b, ++count);
            }
        }

        countMap.forEach((k, v) -> {
            huffmanNodes.add(new HuffmanNode(k, v));
        });

        return huffmanNodes;
    }

    //构建哈夫曼树
    public static HuffmanNode generateHuffmanTree(List<HuffmanNode> nodes) {
        //从小到大排序
        while (nodes.size() > 1) {
            Collections.sort(nodes);

            HuffmanNode left = nodes.remove(0);
            HuffmanNode right = nodes.remove(0);
            HuffmanNode parent = new HuffmanNode(null, (left.weight + right.weight));
            parent.left = left;
            parent.right = right;

            nodes.add(parent);
        }

        return nodes.get(0);
    }

    static Map<Byte, String> huffmanCodes = new HashMap<>();

    //生成哈夫曼编码表
    public static void getHuffmanCodeTable(HuffmanNode node, String code) {
        if (node != null) {
            if (node.data == null) {
                getHuffmanCodeTable(node.left, code + "0");
                getHuffmanCodeTable(node.right, code + "1");
            } else {
                huffmanCodes.put(node.data, code);
            }
        }
    }

    //压缩
    public static byte[] zip(byte[] bytes) {
        StringBuilder builder = new StringBuilder();
        for (byte b : bytes) {
            builder.append(huffmanCodes.get(b));
        }

        int length = (builder.length() + 7) / 8;

        byte[] result = new byte[length];

        int index = 0;
        for (int i = 0; i < builder.length(); i += 8) {
            String str = i + 8 > builder.length() ? builder.substring(i) : builder.substring(i, i + 8);
            result[index] = (byte) Integer.parseInt(str, 2);
            index++;
        }

        return result;
    }

    //解码
    public static byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes) {
        StringBuilder builder = new StringBuilder();
        for (int i = 0; i < huffmanBytes.length; i++) {
            boolean flag = (i == huffmanBytes.length - 1);
            builder.append(byteToBitString(!flag, huffmanBytes[i]));
        }

        List<Byte> list = new ArrayList<>();
        Map<String, Byte> map = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet()) {
            map.put(entry.getValue(), entry.getKey());
        }

        for (int i = 0; i < builder.length(); ) {
            int count = 1;
            Byte b = null;

            while (true) {
                b = map.get(builder.substring(i, i + count));
                if (b == null) {
                    count++;
                } else {
                    break;
                }
            }

            list.add(b);
            i += count;
        }

        byte[] result = new byte[list.size()];
        for (int i = 0; i < list.size(); i++) {
            result[i] = list.get(i);
        }

        return result;
    }

    //byte转二进制
    public static String byteToBitString(boolean flag, byte b) {
        int temp = b;

        if (flag) {
            temp |= 256;
        }
        String str = Integer.toBinaryString(temp);
        if (flag) {
            return str.substring(str.length() - 8);
        }

        return str;
    }
}

class HuffmanNode implements Comparable<HuffmanNode> {
    Byte data;
    int weight;
    HuffmanNode left;
    HuffmanNode right;

    public HuffmanNode(Byte data, int weight) {
        this.data = data;
        this.weight = weight;
    }

    @Override
    public String toString() {
        return "HuffmanNode{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }

    @Override
    public int compareTo(HuffmanNode o) {
        return this.weight - o.weight;
    }

    public void preOrder() {
        preOrder(this);
    }

    private void preOrder(HuffmanNode node) {
        System.out.println(node);

        if (node.left != null) {
            preOrder(node.left);
        }

        if (node.right != null) {
            preOrder(node.right);
        }
    }
}